Update run_pm_dmo_NFW_fixed_timestep.md

from hotwheels_io import *

#

# step 1: config of components

#

# at this stage we do not allocate any resource, we just

# pass the right config parameters to the constructors

# in order to compile the underlying C libraries.

# Step 1: Configure components

# This stage configures components without allocating resources.

# Configurations are passed to constructors to compile the underlying C libraries.

#

#call MPI_Init()

mpi = hwc.MPI().init()

#configure my malloc with 2GB

mym = MyMalloc(alloc_bytes=int(2e9))

#configure the default particle SoA with 1e5 particles

p = SoA(maxpart=int(1e5), mem=mym)

#add the particle soa to the multi-type particle SoAs

soas = SoAs(p, mem=mym)

# configure the timestep class to go from 0 to 1 Gyr

# note: G=43007.1 in units of length=kpc, velocity=km/s, mass = 1e10Msun

# if G=43007.1, masses are in 1e10Msun/h, and radii are in ckpc/h, then this

# constant converts from Gyr to internal units. Not much to comment about it

mpi = hwc.MPI().init() # Initialize MPI

mym = MyMalloc(alloc_bytes=int(2e9)) # Configure memory allocator with 2GB

p = SoA(maxpart=int(1e5), mem=mym) # Configure P to hold 1e5 particles

soas = SoAs(p, mem=mym) # Add P to a multi-type SoA container

# Set up a fixed time-step integrator from 0 to 1 Gyr

# Conversion factor for Gyr to internal units

gyr_to_cu = 3.086e+16 / (1e9 * 3600 * 24 * 365)

ts = integrate.FixedTimeStep(soas, G=43007.1, t_from=0., t_to=1. * gyr_to_cu, MPI=mpi)

ic = NFWIC(rs=100., rho0=1e-6, rs_factor=10.) #init the build of a NFW profile with R<10rs, (note that mass is in units of 1e10Msun)

# init a refined PM grids with 7 stacked PLACEHIRESREGION at smaller and smaller scales

# note that PM needs the time-step integrator class TS to attach its DriftTable and kick callbacks

pm = SuperHiResPM(soas=soas, mem=mym, TS=ts, MPI=mpi, pmgrid=128, grids=8, dt_displacement_factor=0.25)

#configure to compile all modules in the current folder

build = make.Build('./', mpi, pm, ts, mym, *soas.values())

#configure generate SoA headers

headers = OnTheFly(build.build_name, *build.components, generate_user_c=True)

#

# step 2: build SoAs headers and compile C files

#

if mpi.rank == 0: #master rank compile and build headers

ts = integrate.FixedTimeStep(

    soas, 

    G=43007.1,  # Gravitational constant in specific units

    t_from=0., 

    t_to=1. * gyr_to_cu, 

    MPI=mpi

)

# Initialize a NFW profile with scale radius `rs=100` and density `rho0=1e-6`

ic = NFWIC(rs=100., rho0=1e-6, rs_factor=10.)

# Configure a refined PM grid with 7 stacked high-resolution regions

pm = SuperHiResPM( #wrapper to the PM C library

    soas=soas, 

    mem=mym, 

    TS=ts, #will use it to attach gravkick callback

    MPI=mpi, 

    pmgrid=128, 

    grids=8, # number of grids to instantiate

    dt_displacement_factor=0.25 #factor for DtDisplacement

)

build = make.Build('./', mpi, pm, ts, mym, *soas.values()) # Compile all modules in the current directory

headers = OnTheFly(build.build_name, *build.components, generate_user_c=True) # Generate SoA headers

if mpi.rank == 0:  # Master rank handles compilation

    headers.write()

    build.compile()

#

# step 3: allocate resources (e.g. MPI, MyMalloc, and allocations within it)

# Step 2: Allocate resources

#

with (utils.Panic(Build=build) as panic, #attach panic handler to C calls

      utils.Timer(Build=build) as timer, #attach timer handler to C calls

      build.enter(debug=mpi.rank == 0), #parse the compiled object files

      mpi.enter(pm), #attach MPI init info to PM module

      mym.enter(*build.components), #actually allocates the 2GB of ram

      p, #allocate the particle data stracture in the MyMalloc area

      ic.enter(p, mpi.ranks, p.get_maxpart()), #sample the NFW in the particle SoA `p` fields

      pm, #call pm_init() the PM regions

      ts #compute DriftTables if necessary

with (

    utils.Panic(Build=build) as panic,  # Attach panic handler

    utils.Timer(Build=build) as timer,  # Attach timer handler

    build.enter(debug=mpi.rank == 0),  # Parse compiled objects

    mpi.enter(pm),  # Initialize MPI in the PM module

    mym.enter(*build.components),  # Allocate 2GB memory

    p,  # Allocate particle data structure in MyMalloc

    ic.enter(p, mpi.ranks, p.get_maxpart()),  # Sample NFW profile

    pm,  # Initialize PM and compute first accelerations

    ts  # Compute DriftTables if needed

):

        pm.compute_accelerations() #first acc computation

    #

        # step 4: the actual run

    # Step 3: Main simulation loop

    #

    while ts.time < ts.time_end:

        ts.find_timesteps()  # Determine timesteps

        ts.do_first_halfstep_kick()  # First kick (includes drift/kick callbacks)

        ts.drift()  # Update particle positions

        pm.compute_accelerations()  # Recompute accelerations

        ts.do_second_halfstep_kick()  # Second kick

        while ts.time < ts.time_end: #main run.c loop

            ts.find_timesteps() #integrator will find timesteps

            #integrator will call halfstep kick, including PM registered drift and kick callbacks

            ts.do_first_halfstep_kick()

            ts.drift() #integretor will drift

            pm.compute_accelerations() #accelerations

            #integrator will call halfstep kick, including PM registered drift and kick callbacks

            ts.do_second_halfstep_kick()

            if mpi.rank == 0 and steps % 10 == 0: #sometimes, master rank will do a plot

                f, ax = plt.subplots(1)

        # Occasionally, generate plots on the master rank

        if mpi.rank == 0 and ts.steps % 10 == 0:

            fig, ax = plt.subplots(1)

            ax.hist2d(p['pos'][:, 0], p['pos'][:, 1], bins=128)

            ax.set_aspect('equal')

                pat = os.path.join(os.getenv('HW_BUILD', '.'), 'snap%d_rank%d.png' % (steps, mpi.rank))

                f.savefig(pat, bbox_inches='tight', dpi=200)

                plt.close(f)

            path = os.path.join(

                os.getenv('HW_BUILD', '.'), f'snap{ts.steps}_rank{mpi.rank}.png'

            )

            fig.savefig(path, bbox_inches='tight', dpi=200)

            plt.close(fig)

print('simualtion finished')

 No newline at end of file

print('Simulation finished')

 No newline at end of file