Ejemplo n.º 1
0
    def add_diagnostic(self, diagnostic):
        # Call method of parent class
        PICMI_Simulation.add_diagnostic(self, diagnostic)

        # Handle diagnostic
        if diagnostic.step_min is None:
            iteration_min = 0
        else:
            iteration_min = diagnostic.step_min
        if diagnostic.step_max is None:
            iteration_max = np.inf
        else:
            iteration_max = diagnostic.step_max
        # Register field diagnostic
        if type(diagnostic) == PICMI_FieldDiagnostic:
            if diagnostic.data_list is None:
                data_list = ['rho', 'E', 'B', 'J']
            else:
                data_list = diagnostic.data_list
            diag = FieldDiagnostic(period=diagnostic.period,
                                   fldobject=self.fbpic_sim.fld,
                                   comm=self.fbpic_sim.comm,
                                   fieldtypes=data_list,
                                   write_dir=diagnostic.write_dir,
                                   iteration_min=iteration_min,
                                   iteration_max=iteration_max)
        # Register particle diagnostic
        elif type(diagnostic) == PICMI_ParticleDiagnostic:
            species_dict = {}
            for s in diagnostic.species:
                if s.name is None:
                    raise ValueError('When using a species in a diagnostic, '
                                     'its name must be set.')
                species_dict[s.name] = s.fbpic_species
            if diagnostic.data_list is None:
                data_list = ['position', 'momentum', 'weighting']
            else:
                data_list = diagnostic.data_list
            diag = ParticleDiagnostic(period=diagnostic.period,
                                      species=species_dict,
                                      comm=self.fbpic_sim.comm,
                                      particle_data=data_list,
                                      write_dir=diagnostic.write_dir,
                                      iteration_min=iteration_min,
                                      iteration_max=iteration_max)

        # Add it to the FBPIC simulation
        self.fbpic_sim.diags.append(diag)
Ejemplo n.º 2
0
    def add_diagnostic(self, diagnostic):
        # Call method of parent class
        PICMI_Simulation.add_diagnostic(self, diagnostic)

        # Handle iteration_min/max in regular diagnostic
        if type(diagnostic) in [
                PICMI_FieldDiagnostic, PICMI_ParticleDiagnostic
        ]:
            if diagnostic.step_min is None:
                iteration_min = 0
            else:
                iteration_min = diagnostic.step_min
            if diagnostic.step_max is None:
                iteration_max = np.inf
            else:
                iteration_max = diagnostic.step_max

        # Register field diagnostic
        if type(diagnostic) in [
                PICMI_FieldDiagnostic, PICMI_LabFrameFieldDiagnostic
        ]:
            if diagnostic.data_list is None:
                data_list = ['rho', 'E', 'B', 'J']
            else:
                data_list = set()  # Use set to avoid redundancy
                for data in diagnostic.data_list:
                    if data in ['Ex', 'Ey', 'Ez', 'E']:
                        data_list.add('E')
                    elif data in ['Bx', 'By', 'Bz', 'B']:
                        data_list.add('B')
                    elif data in ['Jx', 'Jy', 'Jz', 'J']:
                        data_list.add('J')
                    elif data == 'rho':
                        data_list.add('rho')
                data_list = list(data_list)

        if type(diagnostic) == PICMI_FieldDiagnostic:

            diag = FieldDiagnostic(period=diagnostic.period,
                                   fldobject=self.fbpic_sim.fld,
                                   comm=self.fbpic_sim.comm,
                                   fieldtypes=data_list,
                                   write_dir=diagnostic.write_dir,
                                   iteration_min=iteration_min,
                                   iteration_max=iteration_max)

            # Register particle density diagnostic
            rho_density_list = []
            if diagnostic.data_list is not None:
                for data in diagnostic.data_list:
                    if data.startswith('rho_'):
                        # particle density diagnostics, rho_speciesname
                        rho_density_list.append(data)
            if rho_density_list:
                species_dict = {}
                for data in rho_density_list:
                    sname = data[4:]
                    for s in self.species:
                        if s.name == sname:
                            species_dict[s.name] = s.fbpic_species
                pdd_diag = ParticleChargeDensityDiagnostic(
                    period=diagnostic.period,
                    sim=self.fbpic_sim,
                    species=species_dict,
                    write_dir=diagnostic.write_dir,
                    iteration_min=iteration_min,
                    iteration_max=iteration_max)
                self.fbpic_sim.diags.append(pdd_diag)

        elif type(diagnostic) == PICMI_LabFrameFieldDiagnostic:
            diag = BackTransformedFieldDiagnostic(
                zmin_lab=diagnostic.grid.lower_bound[1],
                zmax_lab=diagnostic.grid.upper_bound[1],
                v_lab=c,
                dt_snapshots_lab=diagnostic.dt_snapshots,
                Ntot_snapshots_lab=diagnostic.num_snapshots,
                gamma_boost=self.gamma_boost,
                period=100,
                fldobject=self.fbpic_sim.fld,
                comm=self.fbpic_sim.comm,
                fieldtypes=diagnostic.data_list,
                write_dir=diagnostic.write_dir)
        # Register particle diagnostic
        elif type(diagnostic) in [
                PICMI_ParticleDiagnostic, PICMI_LabFrameParticleDiagnostic
        ]:
            species_dict = {}
            for s in diagnostic.species:
                if s.name is None:
                    raise ValueError('When using a species in a diagnostic, '
                                     'its name must be set.')
                species_dict[s.name] = s.fbpic_species
            if diagnostic.data_list is None:
                data_list = ['position', 'momentum', 'weighting']
            else:
                data_list = diagnostic.data_list
            if type(diagnostic) == PICMI_ParticleDiagnostic:
                diag = ParticleDiagnostic(period=diagnostic.period,
                                          species=species_dict,
                                          comm=self.fbpic_sim.comm,
                                          particle_data=data_list,
                                          write_dir=diagnostic.write_dir,
                                          iteration_min=iteration_min,
                                          iteration_max=iteration_max)
            else:
                diag = BackTransformedParticleDiagnostic(
                    zmin_lab=diagnostic.grid.lower_bound[1],
                    zmax_lab=diagnostic.grid.upper_bound[1],
                    v_lab=c,
                    dt_snapshots_lab=diagnostic.dt_snapshots,
                    Ntot_snapshots_lab=diagnostic.num_snapshots,
                    gamma_boost=self.gamma_boost,
                    period=100,
                    fldobject=self.fbpic_sim.fld,
                    species=species_dict,
                    comm=self.fbpic_sim.comm,
                    particle_data=data_list,
                    write_dir=diagnostic.write_dir)

        # Add it to the FBPIC simulation
        self.fbpic_sim.diags.append(diag)