How to retrieve virial tensor when running on mutiple MPI ranks?

[Lynn3512138]

Dear Community,

I've been trying to retireve the virial tensor from hoomd-blue because I'm running Brownian dynamics simulation where there is no velocity. When I ran on 1 cpu core and 1 gpu, I wrote a code that works fine as follows:

class virialWriter(hd.custom.Action):
    def __init__(self, filename, mode,flush_interval):
        self.filename = filename
        self.file = h5py.File(filename, mode)
        self.file.attrs["app"] = "hoomd-v5"
        self.file.attrs["version"] = "1.0"

        # Create resizable datasets
        self.virial_ds = self.file.create_dataset(
            "virial", shape=(0, 6), maxshape=(None, 6), dtype='f8', chunks=True
        )
        self.timestep_ds = self.file.create_dataset(
            "timestep", shape=(0,), maxshape=(None,), dtype='i8', chunks=True
        )
        self._cur_frame = 0
        self.flush_interval = flush_interval

    def act(self, timestep):
        #forces = self._state._simulation.operations.integrator.forces
        #all_virials = np.array([f.virials for f in forces])  # (n_forces, n_particles, 6)
        #total_virial = np.sum(all_virials, axis=(0, 1))
        #comm=self._state._simulation.device.communicator.mpi_comm
        with self._state.cpu_local_snapshot as snap:
            virials = snap.particles.net_virial
            #mean_virial = np.mean(virials, axis=0)
            total_virial = np.sum(virials, axis=0)  # shape (6,)
        #global_total_virial=np.zeros_like(local_total_virial)
        #comm.Allreduce(local_total_virial, global_total_virial, op=MPI.SUM)

        # Resize datasets to hold one more row
        self.virial_ds.resize((self._cur_frame + 1, 6))
        self.timestep_ds.resize((self._cur_frame + 1,))

        # Write current frame
        self.virial_ds[self._cur_frame, :] = total_virial
        self.timestep_ds[self._cur_frame] = timestep
        self._cur_frame += 1

        if self._cur_frame % self.flush_interval == 0:
            self.file.flush() 

    def close(self):
        self.file.flush()
        self.file.close()

    def __del__(self):
        self.close()

Then I add it to the operations.writers

vfilename = f"virial_{inputs['fout']}.h5"
if os.path.exists(vfilename):
     os.remove(vfilename)
vi_writer = hd.write.CustomWriter(action=virialWriter(vfilename, "w",
                                                            flush_interval=int(inputs["iofrq_thermo"]/inputs["iofrq_stress"])),
                                        trigger=hd.trigger.Periodic(period=inputs["iofrq_stress"]))
simulation.operations.writers.append(vi_writer)

However, when I try to run the simulation on mutiple cpus, since "cpu_local_snapshot" only returns data on a rank, the virial is wrong. Thus I changed the code to the following

class virialWriter2(hd.custom.Action):
    def __init__(self, forces, filename, mode,flush_interval):
        self.filename = filename
        self.file = h5py.File(filename, mode)
        self.file.attrs["app"] = "hoomd-v5"
        self.file.attrs["version"] = "1.0"
        self.forces=forces

        # Create resizable datasets
        self.virial_ds = self.file.create_dataset(
            "virial", shape=(0, 6), maxshape=(None, 6), dtype='f8', chunks=True
        )
        self.timestep_ds = self.file.create_dataset(
            "timestep", shape=(0,), maxshape=(None,), dtype='i8', chunks=True
        )
        self._cur_frame = 0
        self.flush_interval = flush_interval

    def act(self, timestep):
        #forces = self._state._simulation.operations.integrator.forces
        #constraints = self._state._simulation.operations.integrator.constraints
        total_virial=np.zeros(6)
        for f in self.forces:
            total_virial+=np.sum(f.virials,axis=0)
            total_virial+=f.additional_virial
        #for c in constraints:
        #    total_virial+=np.sum(c.virials,axis=0)
        #    total_virial+=c.additional_virial

        #all_virials = np.array([f.virials for f in forces])  # (n_forces, n_particles, 6)
        #total_virial = np.sum(all_virials, axis=(0, 1))
        
        # Resize datasets to hold one more row
        self.virial_ds.resize((self._cur_frame + 1, 6))
        self.timestep_ds.resize((self._cur_frame + 1,))

        # Write current frame
        self.virial_ds[self._cur_frame, :] = total_virial
        self.timestep_ds[self._cur_frame] = timestep
        self._cur_frame += 1

        if self._cur_frame % self.flush_interval == 0:
            self.file.flush() 

    def close(self):
        self.file.flush()
        self.file.close()

    def __del__(self):
        self.close()

And add it to the writers as follows:

vfilename2 = f"virial2_{inputs['fout']}.h5"
        if os.path.exists(vfilename2):
            os.remove(vfilename2)
        vi_writer2 = hd.write.CustomWriter(action=virialWriter2(filename=vfilename2,forces=integrator.forces, mode="w",
                                                              flush_interval=int(inputs["iofrq_thermo"]/inputs["iofrq_stress"])),
                                          trigger=hd.trigger.Periodic(period=inputs["iofrq_stress"]))
        simulation.operations.writers.append(vi_writer2)

And the same, it works fine using 1 core, but extremely slow when I use 2 cores. Even slower than when I use 1 core.

Therefore, I'd like to ask if anyone has a simple and fast way to output the virial tensor when the simulation is run on mutiple ranks? Thank you very much!

[joaander]

Access the global virial produced by each force directly with the virial attribute of each Force instance: https://hoomd-blue.readthedocs.io/en/v5.4.0/hoomd/md/force/force.html#hoomd.md.force.Force.virials . This method implements a very efficient method to collect and sort the global virials in tag order. It may or may not be faster on 2 cores than 1 - but there is a lot of extra communication needed in the latter.

Also, you should open and write to the file only on rank 0.