Using a callback to modify particle properties

[archermarx]

Hi,

I'm trying to learn about the capabilities of the Python callbacks. I'm trying to use a callback to modify particle properties. As a simple example, say I wanted to set the electron z-momentum to zero each iteration. What I would do is:

# Initialize simulation
sim.initialize_inputs()
sim.initialize_warpx()

elec_wrapper = particle_containers.ParticleContainerWrapper('electrons')
def adjust_velocity():
    nps = elec_wrapper.nps
    uz = elec_wrapper.get_particle_uz()
    uz[:] = np.zeros(len(uz))

callbacks.installbeforestep(adjust_velocity)

# Run simulation for one step
sim.step(1)

However, this doesn't seem to work. Is there a correct way to do this?

[dpgrote]

In WarpX, the particles are divided up into the same blocks that the fields are divided up into. Because of this, there is not a single array holding that particle quantities, but many arrays. The thing that is returned by elec_wrapper.get_particle_uz() is a list of those arrays. To set the z velocity, you would need a loop like this, looping over the list of arrays

for uzi in uz:
    uzi[:] = 0.

[archermarx]

Ah perfect! Thanks!