Esempio n. 1
0
    def propagate_af(self, autocorrelation_function,
                     directory_name="propagation_srw",
                     af_output_file_root=None,
                     maximum_mode=None,
                     python_to_be_used="/users/srio/OASYS1.1/miniconda3/bin/python"):


        propagator = AutocorrelationFunctionPropagator(self._srw_beamline)

        if maximum_mode is None:
            mode_distribution=autocorrelation_function.modeDistribution()
            maximum_mode = mode_distribution[abs(mode_distribution)>0.00005].shape[0]

        propagator.setMaximumMode(maximum_mode)
        data_directory = "%s/%s" % (directory_name, "wavefronts")

        if isMaster():
            if not os.path.exists(directory_name):
                os.mkdir(directory_name)
            if not os.path.exists(data_directory):
                os.mkdir(data_directory)
        barrier()

        propagated_filename = "%s/%s.npz" % (data_directory, "wavefront")
        af = propagator.propagate(autocorrelation_function, propagated_filename,method="SRW",
                                  python_to_be_used=python_to_be_used)

        barrier()
        if isMaster():
            if af_output_file_root is not None:
                af.save("%s.npz" % (af_output_file_root))

        return af
Esempio n. 2
0
    def propagate_af(self, autocorrelation_function,
                     directory_name="propagation_wofry",
                     af_output_file_root=None,
                     maximum_mode=None, # this is indeed the NUMBER OF MODES (maximum index plus one)
                     python_to_be_used="/users/srio/OASYS1.1/miniconda3/bin/python"):

        propagator = AutocorrelationFunctionPropagator(self)

        if maximum_mode is None:
            mode_distribution = autocorrelation_function.modeDistribution()
            maximum_mode = mode_distribution[abs(mode_distribution) > 0.00005].shape[0]

        propagator.setMaximumMode(maximum_mode)
        data_directory = "%s/%s" % (directory_name, "wavefronts")

        if isMaster():
            if not os.path.exists(directory_name):
                os.mkdir(directory_name)
            if not os.path.exists(data_directory):
                os.mkdir(data_directory)
        barrier()

        # propagated_filename = "%s/%s.npz" % (data_directory, af_name)
        propagated_filename = "%s/%s.npz" % (data_directory, "wavefront")
        af = propagator.propagate(autocorrelation_function, propagated_filename, method="WOFRY",
                                  python_to_be_used=python_to_be_used)

        barrier()
        if isMaster():
            if af_output_file_root is not None:
                af.save("%s.npz" % (af_output_file_root))
                logAll("File written to disk: %s.npz" % (af_output_file_root))

        return af
Esempio n. 3
0
    def _estimateMemoryConsumption(self, wavefront):
        configuration = self.configuration()

        size_x = wavefront.absolute_x_coordinates().size
        size_y = wavefront.absolute_y_coordinates().size
        size_wavefront = size_x * size_y
        size_matrix = size_wavefront**2
        size_matrix_per_core = size_matrix / mpi.COMM_WORLD.size
        size_modes = size_wavefront * configuration.numberModes()
        size_solver_worst = size_wavefront * configuration.numberModes() * 12
        size_solver_worst_per_core = size_solver_worst / mpi.COMM_WORLD.size

        if isMaster():
            log("Using sampling factor: %.02f" %
                configuration.samplingFactor())
            log("Using %i x %i points per plane" % (size_x, size_y))
            log("A single wavefront needs %.2f mb" %
                (size_wavefront * 16 / 1024.**2))
            log("The total matrix needs %.2f gb" %
                (size_matrix * 16 / 1024.**3))
            log("Per core the matrix needs %.2f gb" %
                (size_matrix_per_core * 16 / 1024.**3))
            log("The modes will take %.2f mb" % (size_modes * 16 / 1024.**2))
            log("The solver needs in the worst case about %.2f gb" %
                (size_solver_worst * 16 / 1024.**3))
            log("Per core the solver needs in the worst case about %.2f gb" %
                (size_solver_worst_per_core * 16 / 1024.**3))

        return size_matrix
Esempio n. 4
0
    def save(self, filename):
        file_array_shape = self._eigenvectors.shape

        if isMaster():
            fp = np.memmap(filename+".npy", dtype=np.complex128, mode='w+', shape=file_array_shape)

            for i_vector in range(self.numberVectors()):
                logProgress(self.numberVectors(), i_vector, "Writing vectors")

                vector = self.vector(i_vector)

                if isMaster():
                    fp[i_vector, :, :] = vector
            log("Flushing")
            if isMaster():
                del fp
            log("done")
Esempio n. 5
0
    def saveh5(self, filename, af, maximum_number_of_modes=None):
        if has_h5py == False:
            raise ImportError("h5py not available")

        if maximum_number_of_modes is None:
            maximum_number_of_modes = af.Twoform().numberVectors() # af.numberModes()


        if maximum_number_of_modes > af.numberModes():
            maximum_number_of_modes = af.numberModes()

        file_array_shape = (maximum_number_of_modes,
                            len(af.Twoform().xCoordinates()),
                            len(af.Twoform().yCoordinates()))

        if isMaster():
            f = h5py.File(filename, 'w')
            bigdataset = f.create_dataset("twoform_4",file_array_shape, dtype=np.complex)

        for i_vector in range(maximum_number_of_modes):
            logProgress(af.Twoform().numberVectors(), i_vector, "Writing vectors")

            vector = af.Twoform().vector(i_vector)

            if isMaster():
                bigdataset[i_vector,:,:] = vector

        if isMaster():
            sys.stdout.flush()

            data_dict = af.asDictionary()

            for key in data_dict.keys():

                if (key =="twoform_4"): # alreary done
                    pass
                elif (key == "twoform_3"):
                    if (data_dict[key] is not None):
                        f[key] = (data_dict[key])[0:maximum_number_of_modes]
                else:
                    if (data_dict[key] is not None):
                        f[key] = data_dict[key]

            f.close()
Esempio n. 6
0
    def propagate_af(
        self,
        autocorrelation_function,
        directory_name="propagation_wofry",
        af_output_file_root=None,
        maximum_mode=None,  # this is indeed the NUMBER OF MODES (maximum index plus one)
        python_to_be_used="/users/srio/OASYS1.1/miniconda3/bin/python"):

        propagator = AutocorrelationFunctionPropagator(self)

        if maximum_mode is None:
            mode_distribution = autocorrelation_function.modeDistribution()
            maximum_mode = mode_distribution[
                abs(mode_distribution) > 0.00005].shape[0]

        propagator.setMaximumMode(maximum_mode)
        data_directory = "%s/%s" % (directory_name, "wavefronts")

        if isMaster():
            if not os.path.exists(directory_name):
                os.mkdir(directory_name)
            if not os.path.exists(data_directory):
                os.mkdir(data_directory)
        barrier()

        # propagated_filename = "%s/%s.npz" % (data_directory, af_name)
        propagated_filename = "%s/%s.npz" % (data_directory, "wavefront")
        af = propagator.propagate(autocorrelation_function,
                                  propagated_filename,
                                  method="WOFRY",
                                  python_to_be_used=python_to_be_used)

        barrier()
        if isMaster():
            if af_output_file_root is not None:
                af.save("%s.npz" % (af_output_file_root))
                logAll("File written to disk: %s.npz" % (af_output_file_root))

        return af
Esempio n. 7
0
def propagateWavefront(srw_beamline,
                       wavefront,
                       rx,
                       drx,
                       ry,
                       dry,
                       rescale_x,
                       rescale_y,
                       i_mode,
                       python_to_be_used="python"):
    if isMaster():
        if not os.path.exists("tmp"):
            os.mkdir("tmp")

    s_id = str(mpi.COMM_WORLD.Get_rank()) + "_" + gethostname()
    wavefront.save("./tmp/tmp%s_in" % s_id)

    parameter_lines = "rx=%f\ndrx=%f\nry=%f\ndry=%f\nrescale_x=%f\nrescale_y=%f\ns_id=\"%s\"" % (
        rx, drx, ry, dry, rescale_x, rescale_y, s_id)
    pickle.dump(srw_beamline, open("./tmp/tmp%s_beamline.p" % s_id, "wb"))
    lines = """
import pickle
from srwlib import *
from comsyl.waveoptics.SRWAdapter import SRWAdapter
from comsyl.waveoptics.Wavefront import NumpyWavefront, SRWWavefront

wavefront = NumpyWavefront.load("./tmp/tmp%s_in.npz"%s_id)
adapter = SRWAdapter()
wfr = adapter.SRWWavefrontFromWavefront(wavefront,
                                        rx,
                                        drx,
                                        ry,
                                        dry,rescale_x,rescale_y)
#print("Doing propagation in external call")
srw_beamline = pickle.load(open("./tmp/tmp%s_beamline.p"%s_id,"rb"))
srwl.PropagElecField(wfr, srw_beamline)

tmp = SRWWavefront(wfr).toNumpyWavefront()
wfr = None
tmp.save("./tmp/tmp%s_out" % s_id)
"""

    file = open("./tmp/tmp%s.py" % s_id, "w")
    file.writelines(parameter_lines)
    file.writelines("\ni_mode=%d\n" % i_mode)  # added srio
    file.writelines(lines)
    file.close()

    os.system(python_to_be_used + " ./tmp/tmp%s.py" % s_id)

    return NumpyWavefront.load("./tmp/tmp%s_out.npz" % s_id)
Esempio n. 8
0
    def propagate_af(
            self,
            autocorrelation_function,
            directory_name="propagation_srw",
            af_output_file_root=None,
            maximum_mode=None,
            python_to_be_used="/users/srio/OASYS1.1/miniconda3/bin/python"):

        propagator = AutocorrelationFunctionPropagator(self._srw_beamline)

        if maximum_mode is None:
            mode_distribution = autocorrelation_function.modeDistribution()
            maximum_mode = mode_distribution[
                abs(mode_distribution) > 0.00005].shape[0]

        propagator.setMaximumMode(maximum_mode)
        data_directory = "%s/%s" % (directory_name, "wavefronts")

        if isMaster():
            if not os.path.exists(directory_name):
                os.mkdir(directory_name)
            if not os.path.exists(data_directory):
                os.mkdir(data_directory)
        barrier()

        propagated_filename = "%s/%s.npz" % (data_directory, "wavefront")
        af = propagator.propagate(autocorrelation_function,
                                  propagated_filename,
                                  method="SRW",
                                  python_to_be_used=python_to_be_used)

        barrier()
        if isMaster():
            if af_output_file_root is not None:
                af.save("%s.npz" % (af_output_file_root))

        return af
Esempio n. 9
0
    def save(self, filename, af):
        af.Twoform().saveVectors(filename.replace(".npz", ""))

        if isMaster():
            sys.stdout.flush()

            data_dict = af.asDictionary()
            save_dict = dict()
            for key in data_dict.keys():
                if key !="twoform_4":
                    save_dict["np_"+key] = data_dict[key]
                else:
                    eigenvectors = data_dict["twoform_4"]

            filename_npz = filename.replace(".npz", "")

            np.savez_compressed(filename_npz, **save_dict)
def propagateWavefront(srw_beamline, wavefront, rx, drx, ry, dry, rescale_x, rescale_y, i_mode, python_to_be_used="python"):
    if isMaster():
        if not os.path.exists("tmp"):
            os.mkdir("tmp")

    s_id=str(mpi.COMM_WORLD.Get_rank())+"_"+gethostname()
    wavefront.save("./tmp/tmp%s_in"%s_id)

    parameter_lines = "rx=%f\ndrx=%f\nry=%f\ndry=%f\nrescale_x=%f\nrescale_y=%f\ns_id=\"%s\"" % (rx, drx, ry, dry, rescale_x, rescale_y, s_id)
    pickle.dump(srw_beamline, open("./tmp/tmp%s_beamline.p"%s_id,"wb"))
    lines ="""
import pickle
from srwlib import *
from comsyl.waveoptics.SRWAdapter import SRWAdapter
from comsyl.waveoptics.Wavefront import NumpyWavefront, SRWWavefront

wavefront = NumpyWavefront.load("./tmp/tmp%s_in.npz"%s_id)
adapter = SRWAdapter()
wfr = adapter.SRWWavefrontFromWavefront(wavefront,
                                        rx,
                                        drx,
                                        ry,
                                        dry,rescale_x,rescale_y)
#print("Doing propagation in external call")
srw_beamline = pickle.load(open("./tmp/tmp%s_beamline.p"%s_id,"rb"))
srwl.PropagElecField(wfr, srw_beamline)

tmp = SRWWavefront(wfr).toNumpyWavefront()
wfr = None
tmp.save("./tmp/tmp%s_out" % s_id)
"""

    file = open("./tmp/tmp%s.py"%s_id, "w")
    file.writelines(parameter_lines)
    file.writelines("\ni_mode=%d\n"%i_mode) # added srio
    file.writelines(lines)
    file.close()

    os.system(python_to_be_used+" ./tmp/tmp%s.py"%s_id)

    return NumpyWavefront.load("./tmp/tmp%s_out.npz"%s_id)
    def _estimateMemoryConsumption(self, wavefront):
        configuration = self.configuration()

        size_x = wavefront.absolute_x_coordinates().size
        size_y = wavefront.absolute_y_coordinates().size
        size_wavefront = size_x * size_y
        size_matrix = size_wavefront ** 2
        size_matrix_per_core = size_matrix / mpi.COMM_WORLD.size
        size_modes = size_wavefront * configuration.numberModes()
        size_solver_worst = size_wavefront * configuration.numberModes() * 12
        size_solver_worst_per_core = size_solver_worst / mpi.COMM_WORLD.size

        if isMaster():
            log("Using sampling factor: %.02f" % configuration.samplingFactor())
            log("Using %i x %i points per plane" % (size_x, size_y))
            log("A single wavefront needs %.2f mb" % (size_wavefront * 16 / 1024.**2))
            log("The total matrix needs %.2f gb" % (size_matrix * 16 / 1024.**3))
            log("Per core the matrix needs %.2f gb" % (size_matrix_per_core * 16 / 1024.**3))
            log("The modes will take %.2f mb" % (size_modes * 16 / 1024.**2))
            log("The solver needs in the worst case about %.2f gb" % (size_solver_worst * 16 / 1024.**3))
            log("Per core the solver needs in the worst case about %.2f gb" % (size_solver_worst_per_core * 16 / 1024.**3))

        return size_matrix
    def propagate(self, autocorrelation_function, filename, method='SRW', python_to_be_used="python"):

        source_filename = autocorrelation_function._io.fromFile()

        try:
            source_uid = autocorrelation_function.info().uid()
        except:
            source_uid = "None"

        autocorrelation_function.info().logStart()

        logAll("Propagating %s (%s)" % (source_filename, source_uid))


        if self._maximum_mode is None:
            number_modes = autocorrelation_function.numberModes()
        else:
            number_modes = self._maximum_mode

        if isMaster():
            if not os.path.exists("tmp"):
                os.mkdir("tmp")

        distribution_plan = DistributionPlan(mpi.COMM_WORLD, n_rows=number_modes, n_columns=1)

        n_rank = mpi.COMM_WORLD.Get_rank()
        x_coordinates = []
        y_coordinates = []
        for i_mode in distribution_plan.localRows():

            for i in range(1):
                logAll("%i doing mode index: %i/%i (max mode index: %i)" % (n_rank, i_mode, max(distribution_plan.localRows()), number_modes-1))
                if n_rank == 0:
                    sys.stdout.flush()

                wavefront = autocorrelation_function.coherentModeAsWavefront(i_mode)
                #wavefront._e_field[np.abs(wavefront._e_field)<0.000001]=0.0

                if method == 'SRW':
                    # CHANGE THIS FOR WOFRY
                    srw_wavefront = propagateWavefront(self.__srw_beamline,
                                                      wavefront,
                                                      autocorrelation_function.SRWWavefrontRx(),
                                                      autocorrelation_function.SRWWavefrontDRx(),
                                                      autocorrelation_function.SRWWavefrontRy(),
                                                      autocorrelation_function.SRWWavefrontDRy(), 1.0, 1.0, i_mode,
                                                      python_to_be_used=python_to_be_used)
                elif method == 'WOFRY':
                    srw_wavefront = propagateWavefrontWofry(self.__srw_beamline,wavefront,i_mode,python_to_be_used=python_to_be_used)
                else:
                    raise Exception("Method not known: %s"%method)

                # norm_mode = trapez2D( np.abs(srw_wavefront.E_field_as_numpy()[0,:,:,0])**2, 1, 1)**0.5
                # if norm_mode > 1e2 or np.isnan(norm_mode):
                #     print("TRY %i AFTER PROPAGATION:" % i, i_mode,norm_mode)
                #     sys.stdout.flush()
                #else:
                #    break

            #if i==19:
            #    exit()

            # if np.any(norm_srw_wavefront > 10):
            #     exit()
            #
            # if np.any(norm_wavefront > 10):
            #     exit()

            adjusted_wavefront = self._adjustWavefrontSize(srw_wavefront)
            # norm_mode = trapez2D( np.abs(adjusted_wavefront.E_field_as_numpy()[0,:,:,0])**2, 1, 1)**0.5
            # if norm_mode > 1e2 or np.isnan(norm_mode):
            #     print("TRY %i AFTER ADJUSTMENT:" % i, i_mode,norm_mode)
            #     sys.stdout.flush()
            #     exit()

            # writes a file for every wavefront
            TwoformVectorsWavefronts.pushWavefront(filename, adjusted_wavefront, index=i_mode)
            #print("Saving wavefront %i" % i_mode)

            x_coordinates.append(adjusted_wavefront.absolute_x_coordinates().copy())
            y_coordinates.append(adjusted_wavefront.absolute_y_coordinates().copy())

        mpi.COMM_WORLD.barrier()

        # replace the wavefronts bu the propagated ones
        af = self._saveAutocorrelation(autocorrelation_function, number_modes, x_coordinates, y_coordinates, filename)

        # convert from one file per wavefront to one big array
        af.Twoform().convertToTwoformVectorsEigenvectors()
        af.info().setEndTime()

        filelist = glob.glob(filename+"*")
        for f in filelist:
            os.remove(f)

        return af
Esempio n. 13
0
    def calculateAutocorrelation(self, electron_beam, undulator, info):

        configuration = self.configuration()

        # electron_beam = ElectronBeam(energy_in_GeV=6.04,
        #                              energy_spread=0,
        #                              average_current=0.2000,
        #                              electrons=1)

        sigma_matrix = SigmaMatrixFromCovariance(
            xx=electron_beam._moment_xx,
            xxp=electron_beam._moment_xxp,
            xpxp=electron_beam._moment_xpxp,
            yy=electron_beam._moment_yy,
            yyp=electron_beam._moment_yyp,
            ypyp=electron_beam._moment_ypyp,
            sigma_dd=electron_beam._energy_spread,
        )

        resonance_energy = int(
            undulator.resonance_energy(electron_beam.gamma(), 0, 0))
        energy = resonance_energy * configuration.detuningParameter()

        if configuration.virtualSourcePosition() == "":
            if sigma_matrix.isAlphaZero():
                source_position = VIRTUAL_SOURCE_CENTER
            else:
                source_position = VIRTUAL_SOURCE_ENTRANCE
        else:
            source_position = configuration.virtualSourcePosition()

        wavefront_builder = WavefrontBuilder(
            undulator=undulator,
            sampling_factor=configuration.samplingFactor(),
            min_dimension_x=configuration.
            exitSlitWavefrontMinimalSizeHorizontal(),
            max_dimension_x=configuration.
            exitSlitWavefrontMaximalSizeHorizontal(),
            min_dimension_y=configuration.exitSlitWavefrontMinimalSizeVertical(
            ),
            max_dimension_y=configuration.exitSlitWavefrontMaximalSizeVertical(
            ),
            energy=energy,
            source_position=source_position)

        # from comsyl.waveoptics.WavefrontBuilderPySRU import WavefrontBuilderPySRU
        # wavefront_builder = WavefrontBuilderPySRU(undulator=undulator,
        #                                      sampling_factor=configuration.samplingFactor(),
        #                                      min_dimension_x=configuration.exitSlitWavefrontMinimalSizeHorizontal(),
        #                                      max_dimension_x=configuration.exitSlitWavefrontMaximalSizeHorizontal(),
        #                                      min_dimension_y=configuration.exitSlitWavefrontMinimalSizeVertical(),
        #                                      max_dimension_y=configuration.exitSlitWavefrontMaximalSizeVertical(),
        #                                      energy=energy)

        info.setSourcePosition(source_position)

        e_0, sigma_e, beam_energies = self._determineBeamEnergies(
            electron_beam, sigma_matrix, configuration.beamEnergies())

        # determineZ(electron_beam, wavefront_builder, sigma_matrix.sigma_x(), sigma_matrix.sigma_x_prime(),
        #                 sigma_matrix.sigma_y(), sigma_matrix.sigma_y_prime())

        sorted_beam_energies = beam_energies[np.abs(beam_energies).argsort()]

        field_x_coordinates = None
        field_y_coordinates = None
        exit_slit_wavefront = None
        first_wavefront = None

        weighted_fields = None
        for i_beam_energy, beam_energy in enumerate(sorted_beam_energies):
            # TDOO: recheck normalization, especially for delta coupled beams.
            if len(sorted_beam_energies) > 1:
                stepwidth_beam = np.diff(beam_energies)[0]
                weight = (1 / (2 * np.pi * sigma_e**2)**0.5) * np.exp(
                    -beam_energy**2 / (2 * sigma_e**2)) * stepwidth_beam
            else:
                weight = 1.0
            electron_beam._energy = e_0 + beam_energy

            log("%i/%i: doing energy: %e with weight %e" %
                (i_beam_energy + 1, len(beam_energies), electron_beam.energy(),
                 weight))

            # Prepare e_field
            if field_x_coordinates is None or field_y_coordinates is None:
                wavefront = wavefront_builder.build(electron_beam,
                                                    xp=0.0,
                                                    yp=0.0,
                                                    z_offset=0.0)

                reference_wavefront = wavefront.toNumpyWavefront()
            else:
                wavefront = wavefront_builder.buildOnGrid(reference_wavefront,
                                                          electron_beam,
                                                          xp=0.0,
                                                          yp=0.0,
                                                          z_offset=0.0)

            try:
                Rx, dRx, Ry, dRy = wavefront.instantRadii()
            except AttributeError:
                Rx, dRx, Ry, dRy = 0, 0, 0, 0

            energy = wavefront.energies()[0]
            wavefront = wavefront.toNumpyWavefront()

            if exit_slit_wavefront is None:
                exit_slit_wavefront = wavefront.clone()

            wavefront = wavefront_builder.createReferenceWavefrontAtVirtualSource(
                Rx, dRx, Ry, dRy, configuration, source_position, wavefront)

            if self.configuration().useGaussianWavefront() == "true":
                gaussian_wavefront_builder = GaussianWavefrontBuilder()
                wavefront = gaussian_wavefront_builder.fromWavefront(
                    wavefront, info)

            if field_x_coordinates is None or field_y_coordinates is None:
                wavefront = wavefront.asCenteredGrid(resample_x=1.0,
                                                     resample_y=1.0)
                field_x_coordinates = np.array(
                    wavefront.absolute_x_coordinates())
                field_y_coordinates = np.array(
                    wavefront.absolute_y_coordinates())
            else:
                wavefront = wavefront.asCenteredGrid(field_x_coordinates,
                                                     field_y_coordinates)

            size_matrix = self._estimateMemoryConsumption(wavefront)

            if self.adjustMemoryConsumption():
                self._performMemoryConsumptionAdjustment(
                    sigma_matrix, undulator, info, size_matrix)
                exit(0)

            if first_wavefront is None:
                first_wavefront = wavefront.clone()

            if weighted_fields is None:
                weighted_fields = np.zeros(
                    (len(sorted_beam_energies), len(field_x_coordinates),
                     len(field_y_coordinates)),
                    dtype=np.complex128)

            weighted_fields[i_beam_energy, :, :] = np.sqrt(
                weight) * wavefront.E_field_as_numpy()[0, :, :, 0].copy()

        log("Broadcasting electrical fields")
        weighted_fields = mpi.COMM_WORLD.bcast(weighted_fields, root=0)

        static_electron_density, work_matrix = self.calculateAutocorrelationForEnergy(
            wavefront, weighted_fields, sigma_matrix)

        electron_beam._energy = e_0

        if isinstance(work_matrix, AutocorrelationOperator):
            log("Setting up eigenmoder")
            eigenmoder = Eigenmoder(field_x_coordinates, field_y_coordinates)
            log("Starting eigenmoder")

            eigenvalues_spatial, eigenvectors_parallel = eigenmoder.eigenmodes(
                work_matrix, work_matrix.numberModes(), do_not_gather=True)

            if isMaster():
                total_spatial_mode_intensity = eigenvalues_spatial.sum(
                ) * work_matrix._builder._density.normalizationConstant(
                ) / np.trapz(
                    np.trapz(work_matrix.trace(), field_y_coordinates),
                    field_x_coordinates)
                info.setTotalSpatialModeIntensity(total_spatial_mode_intensity)
                log("Total spatial mode intensity: %e" %
                    total_spatial_mode_intensity.real)

            if configuration.twoStepDivergenceMethod() == "":
                divergence_method = "accurate"
            else:
                divergence_method = configuration.twoStepDivergenceMethod()

            divergence_action = DivergenceAction(
                x_coordinates=field_x_coordinates,
                y_coordinates=field_y_coordinates,
                intensity=work_matrix.trace(),
                eigenvalues_spatial=eigenvalues_spatial,
                eigenvectors_parallel=eigenvectors_parallel,
                phase_space_density=PhaseSpaceDensity(
                    sigma_matrix,
                    wavefront.wavenumbers()[0]),
                method=divergence_method)

            twoform = divergence_action.apply(
                number_modes=configuration.numberModes())
        elif isinstance(work_matrix, Twoform):
            twoform = work_matrix
        else:
            eigenmoder = Eigenmoder(field_x_coordinates, field_y_coordinates)
            twoform = eigenmoder.eigenmodes(work_matrix,
                                            configuration.numberModes())

        total_beam_energies = e_0 + beam_energies

        info.setEndTime()
        autocorrelation_function = AutocorrelationFunction(
            sigma_matrix=sigma_matrix,
            undulator=undulator,
            detuning_parameter=configuration.detuningParameter(),
            energy=energy,
            electron_beam_energy=electron_beam.energy(),
            wavefront=first_wavefront,
            exit_slit_wavefront=exit_slit_wavefront,
            srw_wavefront_rx=Rx,
            srw_wavefront_drx=dRx,
            srw_wavefront_ry=Ry,
            srw_wavefront_dry=dRy,
            sampling_factor=configuration.samplingFactor(),
            minimal_size=configuration.exitSlitWavefrontMinimalSizeVertical(),
            beam_energies=total_beam_energies,
            weighted_fields=weighted_fields,
            static_electron_density=static_electron_density,
            twoform=twoform,
            info=info)

        return autocorrelation_function
Esempio n. 14
0
import mpi4py.MPI as mpi
from comsyl.utils.Logger import logAll, log
from comsyl.parallel.utils import isMaster, barrier
from socket import gethostname
from comsyl.parallel.DistributionPlan import DistributionPlan
import os

if isMaster():
    print("Hello from master")

if isMaster():
    if not os.path.exists("tmp"):
        os.mkdir("tmp")

logAll("Using LogAll")
log("Using Log")

s_id = str(mpi.COMM_WORLD.Get_rank()) + "_" + gethostname()

print("s_id: ", s_id)
print("str(mpi.COMM_WORLD.Get_rank()): ", str(mpi.COMM_WORLD.Get_rank()))

number_modes = 1000
distribution_plan = DistributionPlan(mpi.COMM_WORLD,
                                     n_rows=number_modes,
                                     n_columns=1)
print(distribution_plan)

f = open("./tmp/TMP%s_in" % s_id, 'w')
f.write(">>>>>>>>>")
f.close
    def calculateAutocorrelation(self, electron_beam, undulator, info):

        configuration = self.configuration()

        # electron_beam = ElectronBeam(energy_in_GeV=6.04,
        #                              energy_spread=0,
        #                              average_current=0.2000,
        #                              electrons=1)

        sigma_matrix = SigmaMatrixFromCovariance(xx   = electron_beam._moment_xx   ,
                                                 xxp  = electron_beam._moment_xxp  ,
                                                 xpxp = electron_beam._moment_xpxp ,
                                                 yy   = electron_beam._moment_yy   ,
                                                 yyp  = electron_beam._moment_yyp  ,
                                                 ypyp = electron_beam._moment_ypyp ,
                                                 sigma_dd = electron_beam._energy_spread,
                                                 )

        resonance_energy = int(undulator.resonance_energy(electron_beam.gamma(), 0, 0))
        energy = resonance_energy*configuration.detuningParameter()

        if configuration.virtualSourcePosition() == "":
            if sigma_matrix.isAlphaZero():
                source_position = VIRTUAL_SOURCE_CENTER
            else:
                source_position = VIRTUAL_SOURCE_ENTRANCE
        else:
            source_position = configuration.virtualSourcePosition()

        wavefront_builder = WavefrontBuilder(undulator=undulator,
                                             sampling_factor=configuration.samplingFactor(),
                                             min_dimension_x=configuration.exitSlitWavefrontMinimalSizeHorizontal(),
                                             max_dimension_x=configuration.exitSlitWavefrontMaximalSizeHorizontal(),
                                             min_dimension_y=configuration.exitSlitWavefrontMinimalSizeVertical(),
                                             max_dimension_y=configuration.exitSlitWavefrontMaximalSizeVertical(),
                                             energy=energy,
                                             source_position=source_position)

        # from comsyl.waveoptics.WavefrontBuilderPySRU import WavefrontBuilderPySRU
        # wavefront_builder = WavefrontBuilderPySRU(undulator=undulator,
        #                                      sampling_factor=configuration.samplingFactor(),
        #                                      min_dimension_x=configuration.exitSlitWavefrontMinimalSizeHorizontal(),
        #                                      max_dimension_x=configuration.exitSlitWavefrontMaximalSizeHorizontal(),
        #                                      min_dimension_y=configuration.exitSlitWavefrontMinimalSizeVertical(),
        #                                      max_dimension_y=configuration.exitSlitWavefrontMaximalSizeVertical(),
        #                                      energy=energy)

        info.setSourcePosition(source_position)

        e_0, sigma_e, beam_energies = self._determineBeamEnergies(electron_beam, sigma_matrix, configuration.beamEnergies())

        # determineZ(electron_beam, wavefront_builder, sigma_matrix.sigma_x(), sigma_matrix.sigma_x_prime(),
        #                 sigma_matrix.sigma_y(), sigma_matrix.sigma_y_prime())

        sorted_beam_energies = beam_energies[np.abs(beam_energies).argsort()]

        field_x_coordinates = None
        field_y_coordinates = None
        exit_slit_wavefront = None
        first_wavefront = None

        weighted_fields = None
        for i_beam_energy, beam_energy in enumerate(sorted_beam_energies):
            # TDOO: recheck normalization, especially for delta coupled beams.
            if len(sorted_beam_energies) > 1:
                stepwidth_beam = np.diff(beam_energies)[0]
                weight = (1/(2*np.pi*sigma_e**2)**0.5) * np.exp(-beam_energy**2/(2*sigma_e**2)) * stepwidth_beam
            else:
                weight = 1.0
            electron_beam._energy = e_0 + beam_energy


            log("%i/%i: doing energy: %e with weight %e" %(i_beam_energy+1, len(beam_energies),
                                                           electron_beam.energy(),  weight))

            # Prepare e_field
            if field_x_coordinates is None or field_y_coordinates is None:
                wavefront = wavefront_builder.build(electron_beam,
                                                    xp=0.0,
                                                    yp=0.0,
                                                    z_offset=0.0)

                reference_wavefront = wavefront.toNumpyWavefront()
            else:
                wavefront = wavefront_builder.buildOnGrid(reference_wavefront,
                                                          electron_beam,
                                                          xp=0.0,
                                                          yp=0.0,
                                                          z_offset=0.0)

            try:
                Rx, dRx, Ry, dRy = wavefront.instantRadii()
            except AttributeError:
                Rx, dRx, Ry, dRy = 0,0,0,0

            energy = wavefront.energies()[0]
            wavefront = wavefront.toNumpyWavefront()

            if exit_slit_wavefront is None:
                exit_slit_wavefront = wavefront.clone()

            wavefront = wavefront_builder.createReferenceWavefrontAtVirtualSource(Rx, dRx, Ry, dRy, configuration, source_position, wavefront)

            if self.configuration().useGaussianWavefront() == "true":
                gaussian_wavefront_builder = GaussianWavefrontBuilder()
                wavefront = gaussian_wavefront_builder.fromWavefront(wavefront, info)

            if field_x_coordinates is None or field_y_coordinates is None:
                wavefront = wavefront.asCenteredGrid(resample_x=1.0,
                                                     resample_y=1.0)
                field_x_coordinates = np.array(wavefront.absolute_x_coordinates())
                field_y_coordinates = np.array(wavefront.absolute_y_coordinates())
            else:
                wavefront = wavefront.asCenteredGrid(field_x_coordinates, field_y_coordinates)

            size_matrix = self._estimateMemoryConsumption(wavefront)

            if self.adjustMemoryConsumption():
                self._performMemoryConsumptionAdjustment(sigma_matrix, undulator, info, size_matrix)
                exit(0)

            if first_wavefront is None:
                first_wavefront = wavefront.clone()

            if weighted_fields is None:
                weighted_fields = np.zeros((len(sorted_beam_energies), len(field_x_coordinates), len(field_y_coordinates)), dtype=np.complex128)

            weighted_fields[i_beam_energy, :, :] = np.sqrt(weight) * wavefront.E_field_as_numpy()[0, :, :, 0].copy()

        log("Broadcasting electrical fields")
        weighted_fields = mpi.COMM_WORLD.bcast(weighted_fields, root=0)

        static_electron_density, work_matrix = self.calculateAutocorrelationForEnergy(wavefront,
                                                                                      weighted_fields,
                                                                                      sigma_matrix)

        electron_beam._energy = e_0

        if isinstance(work_matrix, AutocorrelationOperator):
            log("Setting up eigenmoder")
            eigenmoder = Eigenmoder(field_x_coordinates, field_y_coordinates)
            log("Starting eigenmoder")

            eigenvalues_spatial, eigenvectors_parallel = eigenmoder.eigenmodes(work_matrix, work_matrix.numberModes(), do_not_gather=True)

            if isMaster():
                total_spatial_mode_intensity = eigenvalues_spatial.sum() * work_matrix._builder._density.normalizationConstant() / np.trapz(np.trapz(work_matrix.trace(), field_y_coordinates), field_x_coordinates)
                info.setTotalSpatialModeIntensity(total_spatial_mode_intensity)
                log("Total spatial mode intensity: %e" % total_spatial_mode_intensity.real)


            if configuration.twoStepDivergenceMethod() == "":
                divergence_method = "accurate"
            else:
                divergence_method = configuration.twoStepDivergenceMethod()

            divergence_action = DivergenceAction(x_coordinates=field_x_coordinates,
                                                 y_coordinates=field_y_coordinates,
                                                 intensity=work_matrix.trace(),
                                                 eigenvalues_spatial=eigenvalues_spatial,
                                                 eigenvectors_parallel=eigenvectors_parallel,
                                                 phase_space_density=PhaseSpaceDensity(sigma_matrix, wavefront.wavenumbers()[0]),
                                                 method=divergence_method)

            twoform = divergence_action.apply(number_modes=configuration.numberModes())
        elif isinstance(work_matrix, Twoform):
            twoform = work_matrix
        else:
            eigenmoder = Eigenmoder(field_x_coordinates, field_y_coordinates)
            twoform = eigenmoder.eigenmodes(work_matrix, configuration.numberModes())

        total_beam_energies = e_0 + beam_energies

        info.setEndTime()
        autocorrelation_function = AutocorrelationFunction(sigma_matrix=sigma_matrix,
                                                           undulator=undulator,
                                                           detuning_parameter=configuration.detuningParameter(),
                                                           energy=energy,
                                                           electron_beam_energy=electron_beam.energy(),
                                                           wavefront=first_wavefront,
                                                           exit_slit_wavefront=exit_slit_wavefront,
                                                           srw_wavefront_rx=Rx,
                                                           srw_wavefront_drx=dRx,
                                                           srw_wavefront_ry=Ry,
                                                           srw_wavefront_dry=dRy,
                                                           sampling_factor=configuration.samplingFactor(),
                                                           minimal_size=configuration.exitSlitWavefrontMinimalSizeVertical(),
                                                           beam_energies=total_beam_energies,
                                                           weighted_fields=weighted_fields,
                                                           static_electron_density=static_electron_density,
                                                           twoform=twoform,
                                                           info=info)

        return autocorrelation_function
Esempio n. 16
0
 def log(self, log_string):
     if isMaster():
         self.logAll(log_string)
Esempio n. 17
0
    def propagate(self,
                  autocorrelation_function,
                  filename,
                  method='SRW',
                  python_to_be_used="python"):

        source_filename = autocorrelation_function._io.fromFile()

        try:
            source_uid = autocorrelation_function.info().uid()
        except:
            source_uid = "None"

        autocorrelation_function.info().logStart()

        logAll("Propagating %s (%s)" % (source_filename, source_uid))

        if self._maximum_mode is None:
            number_modes = autocorrelation_function.numberModes()
        else:
            number_modes = self._maximum_mode

        if isMaster():
            if not os.path.exists("tmp"):
                os.mkdir("tmp")

        distribution_plan = DistributionPlan(mpi.COMM_WORLD,
                                             n_rows=number_modes,
                                             n_columns=1)

        n_rank = mpi.COMM_WORLD.Get_rank()
        x_coordinates = []
        y_coordinates = []
        for i_mode in distribution_plan.localRows():

            for i in range(1):
                logAll("%i doing mode index: %i/%i (max mode index: %i)" %
                       (n_rank, i_mode, max(
                           distribution_plan.localRows()), number_modes - 1))
                if n_rank == 0:
                    sys.stdout.flush()

                wavefront = autocorrelation_function.coherentModeAsWavefront(
                    i_mode)
                #wavefront._e_field[np.abs(wavefront._e_field)<0.000001]=0.0

                if method == 'SRW':
                    # CHANGE THIS FOR WOFRY
                    srw_wavefront = propagateWavefront(
                        self.__srw_beamline,
                        wavefront,
                        autocorrelation_function.SRWWavefrontRx(),
                        autocorrelation_function.SRWWavefrontDRx(),
                        autocorrelation_function.SRWWavefrontRy(),
                        autocorrelation_function.SRWWavefrontDRy(),
                        1.0,
                        1.0,
                        i_mode,
                        python_to_be_used=python_to_be_used)
                elif method == 'WOFRY':
                    srw_wavefront = propagateWavefrontWofry(
                        self.__srw_beamline,
                        wavefront,
                        i_mode,
                        python_to_be_used=python_to_be_used)
                else:
                    raise Exception("Method not known: %s" % method)

                # norm_mode = trapez2D( np.abs(srw_wavefront.E_field_as_numpy()[0,:,:,0])**2, 1, 1)**0.5
                # if norm_mode > 1e2 or np.isnan(norm_mode):
                #     print("TRY %i AFTER PROPAGATION:" % i, i_mode,norm_mode)
                #     sys.stdout.flush()
                #else:
                #    break

            #if i==19:
            #    exit()

            # if np.any(norm_srw_wavefront > 10):
            #     exit()
            #
            # if np.any(norm_wavefront > 10):
            #     exit()

            adjusted_wavefront = self._adjustWavefrontSize(srw_wavefront)
            # norm_mode = trapez2D( np.abs(adjusted_wavefront.E_field_as_numpy()[0,:,:,0])**2, 1, 1)**0.5
            # if norm_mode > 1e2 or np.isnan(norm_mode):
            #     print("TRY %i AFTER ADJUSTMENT:" % i, i_mode,norm_mode)
            #     sys.stdout.flush()
            #     exit()

            # writes a file for every wavefront
            TwoformVectorsWavefronts.pushWavefront(filename,
                                                   adjusted_wavefront,
                                                   index=i_mode)
            #print("Saving wavefront %i" % i_mode)

            x_coordinates.append(
                adjusted_wavefront.absolute_x_coordinates().copy())
            y_coordinates.append(
                adjusted_wavefront.absolute_y_coordinates().copy())

        mpi.COMM_WORLD.barrier()

        # replace the wavefronts bu the propagated ones
        af = self._saveAutocorrelation(autocorrelation_function, number_modes,
                                       x_coordinates, y_coordinates, filename)

        # convert from one file per wavefront to one big array
        af.Twoform().convertToTwoformVectorsEigenvectors()
        af.info().setEndTime()

        filelist = glob.glob(filename + "*")
        for f in filelist:
            os.remove(f)

        return af