Is the field initialization parallelized?

[zhiping0913]

Hi,
Now I am trying to initialize the fields with a given array (which is obtained from other simulations or calculations) by using

Electric_Field_Ey=...   #2D array with shape (Nx,Ny)
Magnetic_Field_Bz=...   #2D array with shape (Nx,Ny)
def initialize_Ey(pos:mp.Vector3):
    x=pos.x+x_center
    y=pos.y+y_center
    Ey=Electric_Field_Ey[get_axis_id_int(axis=x, pos=x), get_axis_id_int(axis=y, pos=y)]
    return Ey
def initialize_Bz(pos:mp.Vector3):
    x=pos.x+x_center
    y=pos.y+y_center
    Bz=Magnetic_Field_Bz[get_axis_id_int(axis=x, pos=x), get_axis_id_int(axis=y, pos=y)]
    return Bz*C.speed_of_light   #convert to meep units

sim = mp.Simulation(...)
sim.init_sim()

#sim.fields.initialize_field(mp.Dx, initialize_Ex)
sim.fields.initialize_field(mp.Dy, initialize_Ey)
#sim.fields.initialize_field(mp.Dz, initialize_Ez)
#sim.fields.initialize_field(mp.Hx, initialize_Bx)
#sim.fields.initialize_field(mp.Hy, initialize_By)
sim.fields.initialize_field(mp.Hz, initialize_Bz)

This initialization method works. However, it is very slow!
For example, in this 2D simulation job

cell_x=60, cell_y=100
Nx=2400, Ny=4000

(That's not a big size)
it takes hours for the initialization.
We can use line_profiler to see what is happening

Total time: 29686.1 s
File: /start_2D.py
Function: initialize_Ey at line 63
Line #    Hits     Time      Per Hit     % Time  Line Contents
==============================================================
   63                                              @profile                               
   64                                              def initialize_Ey(pos:mp.Vector3):     
   65  9619209   8322211.4     0.9         0.0         x=pos.x+x_center                   
   66  9619209   4395234.1     0.5         0.0         y=pos.y+y_center                   
   67  9619209    2.97e+10  3083.5         99.9        Ey=…
   68  9619209  12949366.4     1.3         0.0         return Ey     

Function: main at line 111
Line #    Hits     Time      Per Hit     % Time  Line Contents
==============================================================
...
   125     1    2.99e+10 2.99e+10      31.0      sim.fields.initialize_field(mp.Dy,…
...

The sim.fields.initialize_field calls the initialize_Ey function about Nx*Ny=9600000 times (for every point in the simulation cell), with about 3μs for each call.
The initialization process is not parallelized or even vectorized.
Since the type of the variable of the function initialize_Ey is mp.Vector3 (which is required by sim.fields.initialize_field), it is not easy to vectorize the initialize_Ey function with vectorization tools like jax.vmap.

TypeError: Argument 'Vector3<[1. 2. 3.], [4. 5. 6.], [7. 8. 9.]>' of type '<class 'meep.geom.Vector3'>' is not a valid JAX type

The question is, how to initialize the fields in a more efficiently method?
Can we do the similar initialization with c++ meep?

[zhiping0913]

I tried to jax.jit the get_id_int function, and successfully reduced the time for each call.

@jax.jit
def _get_axis_id_int(axis_start,dr,N, pos):
    id_float=(pos - axis_start)/dr
    id_int=jnp.clip(jnp.round(id_float),0,N-1).astype(jnp.int64)
    return id_int
@jax.jit
def get_x_id_int(pos):
    return _get_axis_id_int(axis_start=x_axis[0], dr=dx, N=Nx, pos=pos)
@jax.jit
def get_y_id_int(pos):
    return _get_axis_id_int(axis_start=y_axis[0], dr=dy, N=Ny, pos=pos)
def initialize_Ex(pos:mp.Vector3):
    x=pos.x+x_center
    y=pos.y+y_center
    Ex=Electric_Field_Ex[get_x_id_int(x), get_y_id_int(y)]
    return Ex

It still took minutes for the initialization

Total time: 479.68 s
File: /start_2D.py
Function: initialize_Ex at line 71
Line #    Hits        Time       Per Hit   % Time  Line Contents
==============================================================
    71                                              @profile                               
    72                                              def initialize_Ex(pos:mp.Vector3):     
    73   9619209    3406122.5      0.4      0.7         x=pos.x+x_center                                    
    74   9619209    3357300.6      0.3      0.7         y=pos.y+y_center
    75   9619209  468507783.6     48.7      97.7        Ex=Electric_Field_Ex[get_x_id_int(x), get_y_id_int(y)]                         
    76   9619209    4408977.3      0.5      0.9         return Ex    

Although jax.jit helps in finding the data in the given array, it still cannot vectorize the entire initialization procedure.
Initialization one point by one is still inefficient.