Logging Add Electrostatics - 1

[xinmeng2020]

Add charge information into H5MD under its standard identifier

• TBD ** see reply below Logging Add Electrostatics - 1  #72 (comment)
• in distribute_input able to define a charges_flag
• assume the charge information store in charges

Electrostatics: treat charge type interaction outside the config.n_types

• The config.n_types contains the \chi interactions; as

phi = pm.create("real", value=0.0)
phi_fourier = ...
force_on_grid = ...

• now define the e interaction separately with suffix _q

    _SPACE_DIM = 3 ## dimension; demo; TBR
    charges_flag = 1 ## demo; TBR
    if charges_flag:
        phi_q = pm.create("real", value=0.0) 
        phi_q_fourier = pm.create("complex", value=0.0)     
        elec_field_fourier= [self.pm.create("complex", value=0.0) for _ in range(_SPACE_DIM)] #for force calculation 
        elec_field = [self.pm.create("real", value=0.0) for _ in range(_SPACE_DIM)] #for force calculation 
        elec_potential_field = self.pm.create("complex", value=0.0) # for energy calculation 
        elec_forces = np.zeros(shape=(len(positions), 3), dtype=dtype)
       ##^----- HERE forces defined as N,3; then not need to transpose like in my old protocol
       ##^----- elec_forces[:,_d] = charges * (elec_field[_d].readout(positions, layout=layout_q))
       ##                                  ^----------- NEED to give column index

Domain Decomposition

• note need to add new defined charges and field_q_forces in the domain_decomposition

dd = domain_decomposition(
     positions,
     pm,
     [velocities,
     [indices,
     [bond_forces,
     [angle_forces,
     [field_forces,
     [names,
     [types,   ---> [*args  in **field.py**  **domain_decomposition**]
    molecules=molecules if molecules_flag else None,
    ...
)

• How to add charges and field_q_forces
  1. by default add; if no charges, then replace with None beforehand?
  2. add conditional judgement based on charges_flag ?
  3. put arguments in a list ---> as tuple ---> unpack together. See #Logging Add Electrostatics - 1  #72 (comment) Clean up the function calls to domain_decomposition from main #73
  4. further upon 3; seems work; see discussions in Clean up the function calls to domain_decomposition from main #73

## way i
dd = domain_decomposition(
     positions,
     pm,
     velocities,
     indices,
     bond_forces,
     angle_forces,
     field_forces,
     charges,  ## <------
     elec_forces, ## <-----
     names,
     types,  
    molecules=molecules if molecules_flag else None,
    ...
)

## way ii
if not charges_flag: # same 
    dd = domain_decomposition(...)
    if molecules_flag:
        (...) = dd 
    else: 
        (...) = dd   
else: # ADD charges and field_q_forces in the |--| place
    dd = domain_decomposition(.. |--|.)
    if molecules_flag:
        (.. |--|.) = dd 
    else: 
        (.. |--|.) = dd 

    ############### way iii, prepare for the DD 
    Affiliated_DD_Args_in_Len = 9 ## can also use the append without defining the length ?
    args_in = [None for _ in range(Affiliated_DD_Args_Len)]
    args_in[:7] = [
         velocities,
         indices,
         bond_forces,
         angle_forces,
         field_forces,
         names, 
         types
    ]
    if charges_flag: ## add charge related 
        args_in[7] = charges 
        args_in[8] = elec_forces  
        
    ## args_recv for the (..) = dd
    args_recv = args_in.copy()
    args_recv.insert(0, positions)
    if molecules_flag:
        args_recv.append(bonds)
        args_recv.append(molecules)
    ## convert to tuple
    args_in = tuple(args_in)
    args_recv= tuple(args_recv)
    
    ############### DD 
    if config.domain_decomposition:
        dd = domain_decomposition(
            positions,
            pm,
            *args_in,
            molecules=molecules if molecules_flag else None,
            bonds=bonds if molecules_flag else None,
            verbose=args.verbose,
            comm=comm,
        )
        args_recv = dd ##<----- wrong

    ##### way iv
    args_in = [
         velocities,
         indices,
         bond_forces,
         angle_forces,
         field_forces,
         names, 
         types
    ]
    args_recv = [
         'positions',
         'velocities',
         'indices',
         'bond_forces',
         'angle_forces',
         'field_forces',
         'names', 
         'types'
    ]

    if charges_flag: ## add charge related 
        args_in.append(charges) 
        args_in.append(elec_forces)
        args_recv.append('charges')
        args_recv.append('elec_forces')

    if molecules_flag:
        args_recv.append('bonds')
        args_recv.append('molecules')
    
    ## convert to tuple
    args_in = tuple(args_in)
    
    ## cmd string to excecut the (...) = dd 
    _str_receive_dd =  ','.join(args_recv)
    _cmd_receive_dd = f"({_str_receive_dd }) = dd"
    
    ############### DD 
    if config.domain_decomposition:
        dd = domain_decomposition(
            positions,
            pm,
            *args_in,
            molecules=molecules if molecules_flag else None,
            bonds=bonds if molecules_flag else None,
            verbose=args.verbose,
            comm=comm,
        )
        exec(_cmd_receive_dd ) 

compute charge (_q) related: field, force, energy

• following treat charge type interaction outside the config.n_types , here update the q related terms in a separate block under the charges_flag, e.g. the code below
• via defining a layout_q separately, it is possible to include denser mesh than the \kai interactions
  • if using different mesh, have to give a separate pm object e.g. pm_q = pmesh.ParticleMesh(config.mesh_size_q, BoxSize=config.box_size, dtype="f4", comm=comm)
  • too much for now?

if not args.disable_field:
    layouts = ...
    update_field(...)
    field_energy, kinetic_energy = compute_field_and_kinetic_energy(...)
    compute_field_force(...)
else:
    kinetic_energy = comm.allreduce(...)

## Add Simple Poisson Equation Electrostatics: compute field/force/energy together 
field_q_energy = 0.0 
if charges_flag:
    layout_q = pm.decompose( positions ) 
    ## ^---- possible to filter out the particles without charge via e.g. positions[charges != 0] following positions[types == t])
    field_q_energy = update_field_force_energy_q(...)

if molecules_flag:
    ...

update_field_force_energy_q

• note that there is slightly difference below; use the latter, which keeps the uncharged particles zero force

  • elec_field on gird --> elec_force on grid --> elec_force readout on particle
  • elec_field on grid --> elec_field readout on particle --> elec_force on particle

• code run success commit 228c2b6

• for the connivence of RESPA energy calculation, update_field_force_energy_q, split to

  • update_field_force_q
  • compute_field_energy_q

def update_field_force_energy_q(
    charges,# charge
    phi_q,  # chage density
    phi_q_fourier,   
    elec_field_fourier, #for force calculation 
    elec_field,     
    elec_forces,    
    elec_energy_field, # for energy calculation 
    field_q_energy, # electric energy
    layout_q, #### general terms  
    pm,
    positions,  
    config,
    compute_energy=False,
    comm=MPI.COMM_WORLD,
    ):
    ## basic setup 
    V = np.prod(config.box_size)
    n_mesh_cells = np.prod(np.full(3, config.mesh_size))
    volume_per_cell = V / n_mesh_cells
    
    ## paint  ## pm.paint(positions[types == t], layout=layouts[t], out=phi[t])
    pm.paint(positions, layout=layout_q, mass=charges, out=phi_q) ## 
    ## scale and fft  
    phi_q /= volume_per_cell 
    phi_q.r2c(out=phi_q_fourier)
    
    def phi_transfer_function(k, v): 
        return v * np.exp(-0.5*config.sigma**2*k.normp(p=2, zeromode=1))
    phi_q_fourier.apply(phi_transfer_function, out=phi_q_fourier) 
    ## ^------ use the same gaussian as the \kai interaciton 
    ## ^------ tbr; phi_transfer_funciton by hamiltonian.H ??
    phi_q_fourier.c2r(out=phi_q) ## this phi_q is after applying the smearing function 
    
    ## electric field via solving poisson equation 
    _SPACE_DIM = 3     
    COULK_GMX = 138.935458 
    for _d in np.arange(_SPACE_DIM):    
        def poisson_transfer_function(k, v, d=_d):
            return - 1j * k[_d] * 4.0 * np.pi * COULK_GMX * v / k.normp(p=2,zeromode=1)
            ######return - 1j * k[_d] * 4.0 * np.pi * v /k.normp(p=2) #hymd.py:173: RuntimeWarning: invalid value encountered in true_divide
        phi_q_fourier.apply(poisson_transfer_function, out = elec_field_fourier[_d])
        elec_field_fourier[_d].c2r(out=elec_field[_d])

    ## calculate electric forces on particles  
    for _d in np.arange(_SPACE_DIM):
        elec_forces[:,_d] = charges * (elec_field[_d].readout(positions, layout=layout_q))
        ###^------ here the use the column, as the elec_forces are defined as (N,3) dimension
    
    ## calculate electric energy in Fourier space
    if compute_energy:
        def transfer_energy(k,v):  ### potential field is electric field / (-ik)  --> potential field * q --> 
            return 4.0 * np.pi * COULK_GMX * np.abs(v)**2  / k.normp(p=2,zeromode=1) ## zeromode = 1 needed here?
        phi_q_fourier.apply(transfer_energy,  kind='wavenumber', out=elec_energy_field)
        
        field_q_energy = 0.5 * comm.allreduce(np.sum(elec_energy_field.value))

    return field_q_energy.real

continued in #74

[mortele]

Add charge information into H5MD under its standard identifier

• TBD

In the input we can probably just store it as a np.float32 dataset in the parent / HDF5 group, so the input would have the structure

mortele:~$ h5dump -n input_example.hdf5
HDF5 "input_example.hdf5" {
FILE_CONTENTS {
 group      /
 dataset    /bonds
 dataset    /coordinates
 dataset    /indices
 dataset    /molecules
 dataset    /names
 dataset    /types
 dataset    /velocities
 dataset    /charges         <-----
 }
}

mortele:~$ h5dump -d /charges input_example.hdf5
HDF5 "input_example.hdf5" {
DATASET "/charges" {
   DATATYPE  H5T_IEEE_F64LE
   DATASPACE  SIMPLE { ( N_PARTICLES ) / ( N_PARTICLES ) }
   DATA {
   (0): ...
   (10): ...
   (20): ...
    ...
   }
}
}

In the output, we need to adhere to the H5MD format, so we have to place it in a dataset at (c.f. H5MD specification)

/particles/all/charge

[mortele]

Actually, we need to decide if we store charges as integers or floats. In the units we use, charges are given in terms of e, so it could make sense to store them as integer values. However, maybe we should retain the flexibility to work with fractional partial charges. Not sure if we will need this in the coarse-grained setting; maybe ask Michele for his opinion.

[xinmeng2020]

Yes, and I would like to have float numbers, which will bring much extendability. 😄

[mortele]

• How to add charges and field_q_forces

  1. by default add; if no charges, then replace with None beforehand?
  2. √√√ add conditional judgement based on charges_flag ?

## way i
dd = domain_decomposition(
     positions,
     pm,
     velocities,
     indices,
     bond_forces,
     angle_forces,
     field_forces,
     charges,  ## <------
     field_q_forces, ## <-----
     names,
     types,  
    molecules=molecules if molecules_flag else None,
    ...
)

## way ii
if not charges_flag: # same 
    dd = domain_decomposition(...)
    if molecules_flag:
        (...) = dd 
    else: 
        (...) = dd   
else: # ADD charges and field_q_forces in the |--| place
    dd = domain_decomposition(.. |--|.)
    if molecules_flag:
        (.. |--|.) = dd 
    else: 
        (.. |--|.) = dd 

Easy short term fix

Short term I think you should just set charges_flag to False and create an empty charges array with zeros if no charges are provided in the input. This way you don't have to worry about any extra logic for the domain decomposition, just add it as one of the *args.

In the not-so-distant future

Medium term, I think we should look into refactoring the code for the domain_decomposition(...) calls. I really don't like these if tests with each clause being 10 lines of arguments.

Maybe we can do some smart packing of the arguments beforehand, since the only thing that the domain decomposition actually cares about is the positions and the molecules arrays. Anything else is just passed as an *args bundle and funneled through layout.exchange to pass them to the correct MPI rank.

I'm not entirely sure, but I think we can do something like the following without making any copies of any actual data:

args = [None for _ in range(12)]
args[:10] = [
    positions, molecules, pm, velocities, indices, bond_forces,
    angle_forces, field_forces, names, types,
]
if charges_flag:
    args[10] = charges
if mass_flag:
    args[11] = masses

args = tuple(args[:10 + int(charges_flag) + int(mass_flag)])
dd = domain_decomposition(
    *args, 
    molecules=molecules if molecules_flag else None,
    bonds=bonds if molecules_flag else None,
    verbose=args.verbose,
    comm=comm,
)
(...) = dd

This would make adding a new conditional array to domain decomposition a matter of adding just the if whatever_flag: args[-1] = whatever instead of literally adding 2ⁿ if test clauses, each being more than 10 lines of code.

[xinmeng2020]

When declare a preset args,

• the args should exclude some items e.g., positions, pm, molecules? (consider keep the domain_decomposition function )@mortele
• should be careful about the order.
  Funny note, in the domain_decomposition(), the order is molecules --> bonds
  whereas, in the return of domain_decomposition() and (...) = dd, the order is reversed

[mortele]

Yea, the positions, pm, bonds, and molecules arguments are special and need to be provided in the correct order. bonds and molecules are keyword arguments on the way in, and on the way out they are provided as the penultimate and the last return objects, respectively. Any order is fine for the rest of the input, as long as the output uses the same order as the call to domain_decomposition, i.e.

                dd = domain_decomposition(
                    positions,
                    pm,

                    you,                                        # <-- 
                    can,                                        # <--
                    put,                                        # <--
                    what,                                       # <--
                    ever,                                       # <--
                    here,                                       # <--

                    molecules=molecules if molecules_flag else None,
                    bonds=bonds if molecules_flag else None,
                    verbose=args.verbose,
                    comm=comm,
                )
                (
                    positions,
                    bonds,

                    you,                                        # <-- 
                    can,                                        # <--
                    put,                                        # <--
                    what,                                       # <--
                    ever,                                       # <--
                    here,                                       # <--

                    bonds,
                    molecules,
                ) = dd

[mortele]

This should be all taken care of now.